import shutil, os, json
import os.path as osp
import cv2
import numpy as np
from tool import filesystem, via_tool


# 重命名---成对
def rename_by_pair2(dest_dir):
    img_files = filesystem.get_all_filepath(dest_dir, [".jpg", ".JPG", ".bmp", ".BMP", ".PNG", ".png"])    
    print(len(img_files))

    rename_cnt =  10000
    # length = len(name_list)
    for idx, f in enumerate(img_files):
        txt_file = ".".join(f.split(".")[:-1]) + ".txt"
        if not osp.exists(txt_file):
            print("not exists: ", txt_file )
            continue
        shutil.move(f, osp.join(dest_dir, "{}.{}".format(rename_cnt, f.split(".")[-1])))
        shutil.move(txt_file, osp.join(dest_dir, "{}.txt".format(rename_cnt)))
        rename_cnt+=1


def distribute_yolo_dataset(data_dir, per_img_in_dir=500):

    img_files = filesystem.get_all_filepath(data_dir, [".jpg"])
    dir_cnt = len(img_files) // per_img_in_dir
    dir_cnt = dir_cnt if len(img_files) % per_img_in_dir == 0 else dir_cnt+1
    
    # sort by name
    sort_data = [[int(osp.basename(p).split(".")[0]), p] for p in img_files]
    sort_data = sorted(sort_data, key=lambda x: x[0])
    for dir_idx in range(dir_cnt):
        s = dir_idx * per_img_in_dir
        e = (dir_idx+1) * per_img_in_dir
        save_dir = osp.join(data_dir, str(dir_idx))
        os.makedirs(save_dir, exist_ok=True)
        for idx in range(s, e):
            if idx >= len(sort_data): continue
            ori_path = sort_data[idx][1]
            txt_file = ".".join(ori_path.split(".")[:-1]) + ".txt"
            shutil.move(ori_path, osp.join(save_dir, osp.basename(ori_path)))
            shutil.move(txt_file, osp.join(save_dir, osp.basename(txt_file)))

    # for d in ori_dirs:
    #     shutil.rmtree(d)


def convert_img_to_jpg(dest_dir):
    img_files = filesystem.get_all_filepath(dest_dir, [".jpg", ".JPG", ".bmp", ".BMP", ".PNG", ".png"])    
    for f in img_files:
        if f.endswith(".jpg"): continue

        img = cv2.imread(f)
        save_path = ".".join(f.split(".")[:-1]) + ".jpg"
        cv2.imwrite(save_path, img)
        os.remove(f)


def txt_to_via_json(data_dir, per_img_in_dir=1000, gen_via_name="via_region_data.json"):
    image_dirs = filesystem.get_last_dir(osp.join(data_dir, "images"))


    for image_dir in image_dirs:
        data_dict = dict()

        txt_dir = image_dir.replace("/images/", "/labels/")
        img_paths = filesystem.get_all_filepath(image_dir)
        for img_path in img_paths:
            filename = osp.basename(img_path)
            txt_path = osp.join(txt_dir, filename.replace(".jpg", ".txt"))
            if not osp.exists(txt_path):
                print("not exists, ", txt_path)
                continue
            
            cv_img = cv2.imread(img_path)
            if cv_img is None:
                print("None: ", img_path)
            # h, w = cv_img.shape[:2]

            regions = []
            rf = open(txt_path, "r")
            for line in rf.readlines():
                outs = line.split(",")
                if len(outs) != 8:
                    # print("len(outs) != 8... ", txt_path)                
                    continue
                
                region = dict()
                label = "tsign"
                xs = []
                ys = []
                for idx, v in enumerate(outs):
                    if idx % 2 == 0: xs.append(int(v))
                    else: ys.append(int(v))
                min_x = int(np.min(xs))
                min_y = int(np.min(ys))
                max_x = int(np.max(xs))
                max_y = int(np.max(ys))

                shape_attribute = {"name":"rect",
                                    "x":min_x,
                                    "y":min_y,
                                    "width":max_x - min_x,
                                    "height":max_y - min_y
                                }
                region["shape_attributes"] = shape_attribute
                region["region_attributes"] = {"label": label}
                regions.append(region)

            rf.close()
            file_size = osp.getsize(img_path)
            data_dict[osp.basename(img_path) + str(file_size)] = {"size": file_size,
                    "filename": osp.basename(img_path),
                    "regions": regions,
                    "file_attributes":{}
                }
        
        with open(image_dir + os.sep + gen_via_name,"w") as wf:
            wf.write(json.dumps(data_dict))

if __name__ == "__main__":

    # dest_dir = '/home/xc/work/dataset/det/TsignDet Train Database'
    # rename_by_pair2(dest_dir)

    # dest_dir = '/home/xc/work/dataset/det/TsignDet Train Database'
    # convert_img_to_jpg(dest_dir)

    # data_dir = "/home/xc/work/dataset/det/TsignDet Train Database/label"
    # distribute_yolo_dataset(data_dir, 300)

    data_dir = "/home/xc/work/dataset/det/TsignDet_Database"
    txt_to_via_json(data_dir)

